#include "Mp4Muxer.h"

static tb_char_t m_recPath[256];
static MP4FileHandle hMp4File;
static tb_bool_t  m_bEnableH265;
static tb_bool_t m_bfindVps;
static tb_bool_t m_bfindSps;
static tb_bool_t m_bfindPps;
static tb_int_t m_nWidth;
static tb_int_t m_nHeight;
static tb_int_t m_nFrameRate;
static tb_int_t m_nAudioSampleRate;
static tb_int_t m_nTimeScale;
static MP4TrackId m_videoId;
static MP4TrackId m_audioId;

tb_void_t enableH265();
tb_void_t disableH265();
tb_int_t readOneNaluFromBuf(tb_byte_t *buffer, tb_int_t nBufferSize,
                            tb_int_t offSet, MP4ENC_NaluUnit *nalu);
tb_int_t writeH264Track(tb_byte_t *pframeBuf, tb_int_t frameLen);
tb_int_t writeH265Track(tb_byte_t *pframeBuf, tb_int_t frameLen);
tb_int_t writeAudioTrack(tb_byte_t *pframeBuf, tb_int_t frameLen);

void Mp4MuxerInit(){
  m_bEnableH265 = true;
  m_bfindVps = tb_false;
  m_bfindSps = tb_false;
  m_bfindPps = tb_false;
  m_nWidth = 1920;
  m_nHeight = 1080;
  m_nFrameRate = 20;
  m_nTimeScale = 90000;
  m_nAudioSampleRate = 16000;
  tb_memset(m_recPath, 0, sizeof(m_recPath));
}

tb_int_t readOneNaluFromBuf(tb_byte_t *buffer, tb_int_t nBufferSize,
                                      tb_int_t offSet, MP4ENC_NaluUnit *nalu) {
  tb_int_t i = offSet;
  while (i < nBufferSize) {
    if (buffer[i++] == 0x00 && buffer[i++] == 0x00 && buffer[i++] == 0x00 &&
        buffer[i++] == 0x01) {
      tb_int_t pos = i;
      while (pos < nBufferSize) {
        if (buffer[pos++] == 0x00 && buffer[pos++] == 0x00 &&
            buffer[pos++] == 0x00 && buffer[pos++] == 0x01) {
          break;
        }
      }
      if (pos == nBufferSize) {
        nalu->frameLen = pos - i;
      } else {
        nalu->frameLen = (pos - 4) - i;
      }
      nalu->frameType = buffer[i] >> 1;
      nalu->pframeBuf = (tb_byte_t *)&buffer[i];
      return (nalu->frameLen + i - offSet);
    }
  }
  return 0;
}

tb_int_t writeAudioTrack(tb_byte_t *pframeBuf, tb_int_t frameLen) {
  int ret = MP4WriteSample(hMp4File, m_audioId, pframeBuf, frameLen,
                           MP4_INVALID_DURATION, 0, 1);
  if (!ret) {
//    LOGE("Audio record MP4WriteSample : %d\n", ret);
  }
  return 0;
}

tb_int_t writeH264Track(tb_byte_t *pframeBuf, tb_int_t frameLen) {
  tb_int_t pos = 0, len = 0;
  tb_int_t nalu_type = 0;
  MP4ENC_NaluUnit nalu;
  /* LOGE("\n\nframeLen:%d\n\n", frameLen); */
  while ((len = readOneNaluFromBuf(pframeBuf, frameLen, pos, &nalu))) {
    /* LOGE("\n\nnalu.frameType:0x%02x nalu.frameLen:%d len:%d\n\n",
     * nalu.frameType, nalu.frameLen, len); */
    if (nalu.frameType == 0x07) // sps
    {
      if (m_bfindSps)
        goto continue_pos;

      // 添加h264 track
      m_videoId = MP4AddH265VideoTrack(
          hMp4File, m_nTimeScale, m_nTimeScale / m_nFrameRate,
          m_nWidth,          // width
          m_nHeight,         // height
          nalu.pframeBuf[1], // sps[1] AVCProfileIndication
          nalu.pframeBuf[2], // sps[2] profile_compat
          nalu.pframeBuf[3], // sps[3] AVCLevelIndication
          3);                // 4 bytes length before each NAL unit
      if (m_videoId == MP4_INVALID_TRACK_ID) {
//        LOGE("add video track failed.\n");
        return 0;
      }
      MP4SetVideoProfileLevel(hMp4File, 1); //  Simple Profile @ Level 3
      MP4AddH264SequenceParameterSet(hMp4File, m_videoId, nalu.pframeBuf,
                                     nalu.frameLen);

      m_bfindSps = true;
    } else if (nalu.frameType == 0x08) // pps
    {
      if (m_bfindPps)
        goto continue_pos;

      MP4AddH264PictureParameterSet(hMp4File, m_videoId, nalu.pframeBuf,
                                    nalu.frameLen);
      m_bfindPps = true;
    } else {
      if ((m_videoId == MP4_INVALID_TRACK_ID) || !m_bfindSps || !m_bfindPps) {
        goto continue_pos;
      }

      tb_int_t datalen = nalu.frameLen + 4;
      tb_byte_t *data = (tb_byte_t *)tb_malloc0(datalen);
      // MP4 Nalu前四个字节表示Nalu长度
      data[0] = nalu.frameLen >> 24;
      data[1] = nalu.frameLen >> 16;
      data[2] = nalu.frameLen >> 8;
      data[3] = nalu.frameLen & 0xff;
      nalu_type = nalu.frameType;
      tb_memcpy(data + 4, nalu.pframeBuf, nalu.frameLen);
      // LOGE("debug step 2 ...........................\n");
      int if_I = 0;
      if(nalu_type == 19){
        if_I = 1;
      }
      tb_int_t ret = MP4WriteSample(hMp4File, m_videoId, data, datalen,
                                    MP4_INVALID_DURATION, 0, if_I);
      // LOGE("ret : %d.....\n",ret);   //1正常，0异常
      if (!ret) {
        tb_free(data);
        return -1;
      }
      tb_free(data);
    }
  continue_pos:
    pos += len;
  }
  return nalu_type;
}

tb_int_t writeH265Track(tb_byte_t *pframeBuf, tb_int_t frameLen) {
  tb_int_t pos = 0, len = 0;
  tb_int_t nalu_type = 0;
  MP4ENC_NaluUnit nalu;
//  LOGE("\n\nframeLen:%d\n\n", frameLen);
  while ( (len = readOneNaluFromBuf(pframeBuf, frameLen, pos, &nalu)) ) {
//    LOGE("\n\nnalu.frameType:0x%02x nalu.frameLen:%d len:%d\n\n",
//      nalu.frameType, nalu.frameLen, len); 
    if (nalu.frameType == 32) // vps
    {
      if (m_bfindVps)
        goto continue_pos;

      // 添加h265 track
      m_videoId = MP4AddH265VideoTrack(
          hMp4File, m_nTimeScale, m_nTimeScale / m_nFrameRate,
          m_nWidth,          // width
          m_nHeight,         // height
          nalu.pframeBuf[1], // sps[1] AVCProfileIndication
          nalu.pframeBuf[2], // sps[2] profile_compat
          nalu.pframeBuf[3], // sps[3] AVCLevelIndication
          3);                // 4 bytes length before each NAL unit
      if (m_videoId == MP4_INVALID_TRACK_ID) {
//        LOGE("add video track failed.\n");
        return 0;
      }
      MP4SetVideoProfileLevel(hMp4File, 1); //  Simple Profile @ Level 3
      MP4AddH265VideoParameterSet(hMp4File, m_videoId, nalu.pframeBuf,
                                  nalu.frameLen);
      m_bfindVps = tb_true;
    } else if (nalu.frameType == 33) // sps;
    {
      if (m_bfindSps)
        goto continue_pos;
      MP4AddH265SequenceParameterSet(hMp4File, m_videoId, nalu.pframeBuf,
                                     nalu.frameLen);
      m_bfindSps = tb_true;
    } else if (nalu.frameType == 34) // pps;
    {
      if (m_bfindPps)
        goto continue_pos;
      MP4AddH265PictureParameterSet(hMp4File, m_videoId, nalu.pframeBuf,
                                    nalu.frameLen);
      m_bfindPps = tb_true;
    } else {
      if ((m_videoId == MP4_INVALID_TRACK_ID) || !m_bfindVps || !m_bfindSps ||
          !m_bfindPps) {
        goto continue_pos;
      }
      tb_int_t datalen = nalu.frameLen + 4;
      tb_byte_t *data = (tb_byte_t *)tb_malloc0(datalen+8);//多加8字节防止内存越界
      // MP4 Nalu前四个字节表示Nalu长度
      data[0] = nalu.frameLen >> 24;
      data[1] = nalu.frameLen >> 16;
      data[2] = nalu.frameLen >> 8;
      data[3] = nalu.frameLen & 0xff;
      nalu_type = nalu.frameType;
      tb_memcpy(data + 4, nalu.pframeBuf, nalu.frameLen);
      // LOGE("debug step 2 ...........................\n");
      int if_I = 0;
      if(nalu_type == 19){
        if_I = 1;
      }
      tb_int_t ret = MP4WriteSample(hMp4File, m_videoId, data, datalen,
                                    MP4_INVALID_DURATION, 0, if_I);
      // LOGE("ret : %d.....\n",ret);   //1正常，0异常
      if (!ret) {
        tb_free(data);
        return -1;
      }
      tb_free(data);
    }
  continue_pos:
    pos += len;
  }
  return nalu_type;
}

tb_int_t fileCreate(tb_char_t *pFileName, tb_int_t width,
                              tb_int_t height, tb_int_t frameRate) {
  Mp4MuxerInit();
  if (pFileName == NULL) {
      LOGE("%s is null", pFileName);
      return -1;
  }
  snprintf(m_recPath, sizeof(m_recPath), "%s", pFileName);
  m_nWidth = width;
  m_nHeight = height; 
  m_nFrameRate = frameRate;
  // create mp4 file
  hMp4File = MP4Create(pFileName, 0);
  if (hMp4File == MP4_INVALID_FILE_HANDLE) {
    LOGE("ERROR:Open file fialed. fname=%s\n", pFileName);
    return -1;
  }
  MP4SetTimeScale(hMp4File, m_nTimeScale);

  m_audioId = MP4AddALawAudioTrack(hMp4File, m_nAudioSampleRate, 40);
  MP4SetTrackIntegerProperty(hMp4File, m_audioId, "mdia.minf.stbl.stsd.alaw.channels", 1);
  MP4SetAudioProfileLevel(hMp4File, 0x02);
  LOGE("create %s ok\n", pFileName);
  return 0;
}

tb_int_t fileWrite(tb_int_t frameType, tb_byte_t *pframeBuf,
                             tb_int_t frameLen) {
  if (!hMp4File || !pframeBuf) {
    return -1;
  }
  switch (frameType) {
  case MEDIA_FRAME_AUDIO: {
    return writeAudioTrack(pframeBuf, frameLen);
  }
  case MEDIA_FRAME_VIDEO: {
//      LOGE("george hree, %d\n", m_bEnableH265);
    if (m_bEnableH265)
      return writeH265Track(pframeBuf, frameLen);
    else
      return writeH264Track(pframeBuf, frameLen);
  }
  default:
    break;
  }
  return 0;
}

tb_int_t fileClose() {
  if (hMp4File) {
    MP4Close(hMp4File, 0);
    hMp4File = tb_null;
    m_bfindVps = tb_false;
    m_bfindSps = tb_false;
    m_bfindPps = tb_false;
    m_videoId = 0;
  }
  return 0;
}

tb_void_t enableH265(){
  m_bEnableH265 = tb_true;
}


static tb_int_t get_nalu_type(tb_char_t *buf, tb_int_t len) {
  if (len < 6) {
    return -1;
  }
  if (buf[0] != 0 || buf[1] != 0 || buf[2] != 0 || buf[3] != 1) {
    return -1;
  }
  tb_uchar_t nalu_code = buf[4];
  //先判断是否h265帧类型
  tb_uchar_t nalu_type = ((nalu_code)&0x7E) >> 1;
  if (nalu_type == 32 /*VPS*/ || nalu_type == 19 /*IDR*/) {
    return 1;
  } else if (nalu_type == 1) {
    return 0;
  }
  return -1;
}

//TODO: 参考下面调用方法，尝试合成mp4文件
#if 0
//1.创建对象实例
if (fileCreate(recPath, width, height, frameRate) != 0) {
  LogError("mp4 %s fileCreate error", argv[1]);
  return -1;
}
LogDebug("create mp4 %s success !\n",argv[1]);
//2.读音视频帧

//3.将视频写入mp4文件
fileWrite(MEDIA_FRAME_VIDEO, (tb_byte_t*)video_buf, video_len);
//4.将音频写入mp4文件
fileWrite(MEDIA_FRAME_AUDIO, (tb_byte_t*)audio_buf, audio_len);
//5.关闭文件完成想写入
fileClose();
#endif
